#include "GLSLProgram.h"
#include <fstream>
#include <sstream>
#include <sys/stat.h>
#include <vector>
#include <iostream>
using std::ifstream;
using std::ios;
using std::string;


std::map<std::string, ShaderType> extensions = {
    {".vs",   VERTEX},
    {".vert", VERTEX},
    {"_vert.glsl", VERTEX},
    {".vert.glsl", VERTEX },
    {".gs",   GEOMETRY},
    {".geom", GEOMETRY},
    { ".geom.glsl", GEOMETRY },
    { "_geo.glsl", GEOMETRY },
    {".tcs",  TESS_CONTROL},
    { "_tcs.glsl",  TESS_CONTROL },
    { ".tes",  TESS_EVALUATION },
    {"_tes.glsl",  TESS_EVALUATION},
    {".fs",   FRAGMENT},
    {".frag", FRAGMENT},
    {"_frag.glsl", FRAGMENT},
    {".frag.glsl", FRAGMENT},
    {".cs",   COMPUTE},
    { "_cs.glsl", COMPUTE }
};


GLSLProgram::GLSLProgram() : handle(0), linked(false) {}

GLSLProgram::~GLSLProgram() {
    if (handle == 0) return;
    detachAndDeleteShaderObjects();
    glDeleteProgram(handle);
}

GLint GLSLProgram::getUniformLocation(const char* name)
{
    auto pos = uniformLocations.find(name);
    if (pos == uniformLocations.end()) {
        GLint loc = glGetUniformLocation(handle, name);
        uniformLocations[name] = loc;
        return loc;
    }
    return pos->second;
}

void GLSLProgram::detachAndDeleteShaderObjects() {
    GLint numShaders = 0;
    glGetProgramiv(handle, GL_ATTACHED_SHADERS, &numShaders);
    std::vector<GLuint> shaderNames(numShaders);
    glGetAttachedShaders(handle, numShaders, NULL, shaderNames.data());
    for (GLuint shader : shaderNames) {
        glDetachShader(handle, shader);
        glDeleteShader(shader);
    }
}

void GLSLProgram::compileShader(const char* fileName) {

    string ext = getExtension(fileName);
    ShaderType type = VERTEX;
    auto it = extensions.find(ext);
    if (it != extensions.end()) {
        type = it->second;
    }
    else {
        string msg = "Unrecognized extension: " + ext;
        std::cout << msg <<std::endl;
    }

    compileShader(fileName, type);
}

string GLSLProgram::getExtension(const char* name) {
    string nameStr(name);

    size_t dotLoc = nameStr.find_last_of('.');
    if (dotLoc != string::npos) {
        std::string ext = nameStr.substr(dotLoc);
        if (ext == ".glsl") {

            size_t loc = nameStr.find_last_of('.', dotLoc - 1);
            if (loc == string::npos) {
                loc = nameStr.find_last_of('_', dotLoc - 1);
            }
            if (loc != string::npos) {
                return nameStr.substr(loc);
            }
        }
        else {
            return ext;
        }
    }
    return "";
}

void GLSLProgram::compileShader(const char* fileName,
    ShaderType type) {
    if (!fileExists(fileName)) {
        string message = string("Shader: ") + fileName + " not found.";
        std::cout << message << std::endl;
    }

    if (handle <= 0) {
        handle = glCreateProgram();
        if (handle == 0) {
            std::cout << "Unable to create shader program." << std::endl;
        }
    }

    ifstream inFile(fileName, ios::in);
    if (!inFile) {
        string message = string("Unable to open: ") + fileName;
        std::cout<< message <<std::endl;
    }

    std::stringstream code;
    code << inFile.rdbuf();
    inFile.close();

    compileShader(code.str(), type, fileName);
}

void GLSLProgram::compileShader(const string& source,
    ShaderType type,
    const char* fileName) {
    if (handle <= 0) {
        handle = glCreateProgram();
        if (handle == 0) {
            std::cout << "Unable to create shader program." << std::endl;
        }
    }

    GLuint shaderHandle = glCreateShader(type);

    const char* c_code = source.c_str();
    glShaderSource(shaderHandle, 1, &c_code, NULL);

    // Compile the shader
    glCompileShader(shaderHandle);

    // Check for errors
    int result;
    glGetShaderiv(shaderHandle, GL_COMPILE_STATUS, &result);
    if (GL_FALSE == result) {
        // Compile failed, get log
        std::string msg;
        if (fileName) {
            msg = string(fileName) + ": shader compliation failed\n";
        }
        else {
            msg = "Shader compilation failed.\n";
        }

        int length = 0;
        glGetShaderiv(shaderHandle, GL_INFO_LOG_LENGTH, &length);
        if (length > 0) {
            std::string log(length, ' ');
            int written = 0;
            glGetShaderInfoLog(shaderHandle, length, &written, &log[0]);
            msg += log;
        }
        std::cout << msg << std::endl;
    }
    else {
        // Compile succeeded, attach shader
        glAttachShader(handle, shaderHandle);
    }
}

void GLSLProgram::link() {
    if (linked) return;
    if (handle <= 0) {
        std::cout << "Program has not been compiled." << std::endl;
    }

    glLinkProgram(handle);
    int status = 0;
    std::string errString;
    glGetProgramiv(handle, GL_LINK_STATUS, &status);
    if (GL_FALSE == status) {
        // Store log and return false
        int length = 0;
        glGetProgramiv(handle, GL_INFO_LOG_LENGTH, &length);
        errString += "Program link failed:\n";
        if (length > 0) {
            std::string log(length, ' ');
            int written = 0;
            glGetProgramInfoLog(handle, length, &written, &log[0]);
            errString += log;
        }
    }
    else {
        findUniformLocations();
        linked = true;
    }

    detachAndDeleteShaderObjects();

    if (GL_FALSE == status) {
        std::cout << errString << std::endl;
    }
}

void GLSLProgram::findUniformLocations() {
    uniformLocations.clear();

    GLint numUniforms = 0;
    glGetProgramInterfaceiv(handle, GL_UNIFORM, GL_ACTIVE_RESOURCES, &numUniforms);

    GLenum properties[] = { GL_NAME_LENGTH, GL_TYPE, GL_LOCATION, GL_BLOCK_INDEX };

    for (GLint i = 0; i < numUniforms; ++i) {
        GLint results[4];
        glGetProgramResourceiv(handle, GL_UNIFORM, i, 4, properties, 4, NULL, results);

        if (results[3] != -1) continue;  // Skip uniforms in blocks
        GLint nameBufSize = results[0] + 1;
        char* name = new char[nameBufSize];
        glGetProgramResourceName(handle, GL_UNIFORM, i, nameBufSize, NULL, name);
        uniformLocations[name] = results[2];
        delete[] name;
    }
}

void GLSLProgram::use() {
    if (handle <= 0 || (!linked)) {
        std::cout << "Shader has not been linked" << std::endl;
    }
    glUseProgram(handle);
}

int GLSLProgram::getHandle() {
    return handle;
}

bool GLSLProgram::isLinked() {
    return linked;
}

void GLSLProgram::bindAttribLocation(GLuint location, const char* name) {
    glBindAttribLocation(handle, location, name);
}

void GLSLProgram::bindFragDataLocation(GLuint location, const char* name) {
    glBindFragDataLocation(handle, location, name);
}

void GLSLProgram::setUniform(const char* name, float x, float y, float z) {
    GLint loc = getUniformLocation(name);
    glUniform3f(loc, x, y, z);
}

void GLSLProgram::setUniform(const char* name, const glm::vec3& v) {
    this->setUniform(name, v.x, v.y, v.z);
}

void GLSLProgram::setUniform(const char* name, const glm::vec4& v) {
    GLint loc = getUniformLocation(name);
    glUniform4f(loc, v.x, v.y, v.z, v.w);
}

void GLSLProgram::setUniform(const char* name, const glm::vec2& v) {
    GLint loc = getUniformLocation(name);
    glUniform2f(loc, v.x, v.y);
}

void GLSLProgram::setUniform(const char* name, const glm::mat4& m) {
    GLint loc = getUniformLocation(name);
    glUniformMatrix4fv(loc, 1, GL_FALSE, &m[0][0]);
}

void GLSLProgram::setUniform(const char* name, const glm::mat3& m) {
    GLint loc = getUniformLocation(name);
    glUniformMatrix3fv(loc, 1, GL_FALSE, &m[0][0]);
}

void GLSLProgram::setUniform(const char* name, float val) {
    GLint loc = getUniformLocation(name);
    glUniform1f(loc, val);
}

void GLSLProgram::setUniform(const char* name, int val) {
    GLint loc = getUniformLocation(name);
    glUniform1i(loc, val);
}

void GLSLProgram::setUniform(const char* name, GLuint val) {
    GLint loc = getUniformLocation(name);
    glUniform1ui(loc, val);
}

void GLSLProgram::setUniform(const char* name, bool val) {
    int loc = getUniformLocation(name);
    glUniform1i(loc, val);
}

void GLSLProgram::printActiveUniforms() {

    // For OpenGL 4.3 and above, use glGetProgramResource
    GLint numUniforms = 0;
    glGetProgramInterfaceiv(handle, GL_UNIFORM, GL_ACTIVE_RESOURCES, &numUniforms);

    GLenum properties[] = { GL_NAME_LENGTH, GL_TYPE, GL_LOCATION, GL_BLOCK_INDEX };

    std::cout << "************ Active uniforms ************"<<std::endl;
    for (int i = 0; i < numUniforms; ++i) {
        GLint results[4];
        glGetProgramResourceiv(handle, GL_UNIFORM, i, 4, properties, 4, NULL, results);

        if (results[3] != -1) continue;  // Skip uniforms in blocks
        GLint nameBufSize = results[0] + 1;
        char* name = new char[nameBufSize];
        glGetProgramResourceName(handle, GL_UNIFORM, i, nameBufSize, NULL, name);
        std::cout << "* ID:"<< results[2]<< " Name:"<< name<<" Type:"<<getTypeString(results[1])<<std::endl;
        delete[] name;
    }
    std::cout << "*****************************************" << std::endl;
}

void GLSLProgram::printActiveUniformBlocks() {

    GLint numBlocks = 0;

    glGetProgramInterfaceiv(handle, GL_UNIFORM_BLOCK, GL_ACTIVE_RESOURCES, &numBlocks);
    GLenum blockProps[] = { GL_NUM_ACTIVE_VARIABLES, GL_NAME_LENGTH };
    GLenum blockIndex[] = { GL_ACTIVE_VARIABLES };
    GLenum props[] = { GL_NAME_LENGTH, GL_TYPE, GL_BLOCK_INDEX };

    for (int block = 0; block < numBlocks; ++block) {
        GLint blockInfo[2];
        glGetProgramResourceiv(handle, GL_UNIFORM_BLOCK, block, 2, blockProps, 2, NULL, blockInfo);
        GLint numUnis = blockInfo[0];

        char* blockName = new char[blockInfo[1] + 1];
        glGetProgramResourceName(handle, GL_UNIFORM_BLOCK, block, blockInfo[1] + 1, NULL, blockName);
        std::cout << "Uniform block :" << blockName << std::endl;
        delete[] blockName;

        GLint* unifIndexes = new GLint[numUnis];
        glGetProgramResourceiv(handle, GL_UNIFORM_BLOCK, block, 1, blockIndex, numUnis, NULL, unifIndexes);

        for (int unif = 0; unif < numUnis; ++unif) {
            GLint uniIndex = unifIndexes[unif];
            GLint results[3];
            glGetProgramResourceiv(handle, GL_UNIFORM, uniIndex, 3, props, 3, NULL, results);

            GLint nameBufSize = results[0] + 1;
            char* name = new char[nameBufSize];
            glGetProgramResourceName(handle, GL_UNIFORM, uniIndex, nameBufSize, NULL, name);
            std::cout << "   " << name << "    " << getTypeString(results[1]) << std::endl;
            delete[] name;
        }

        delete[] unifIndexes;
    }
}

void GLSLProgram::printActiveAttribs() {

    // >= OpenGL 4.3, use glGetProgramResource
    GLint numAttribs;
    glGetProgramInterfaceiv(handle, GL_PROGRAM_INPUT, GL_ACTIVE_RESOURCES, &numAttribs);

    GLenum properties[] = { GL_NAME_LENGTH, GL_TYPE, GL_LOCATION };

    std::cout << "************ Active attributes ************"<<std::endl;
    for (int i = 0; i < numAttribs; ++i) {
        GLint results[3];
        glGetProgramResourceiv(handle, GL_PROGRAM_INPUT, i, 3, properties, 3, NULL, results);

        GLint nameBufSize = results[0] + 1;
        char* name = new char[nameBufSize];
        glGetProgramResourceName(handle, GL_PROGRAM_INPUT, i, nameBufSize, NULL, name);
        std::cout << "* ID:" << results[2] << " Name:" << name << " Type:" << getTypeString(results[1]) << std::endl;
        delete[] name;
    }
    std::cout << "*******************************************"<<std::endl;
}

const char* GLSLProgram::getTypeString(GLenum type) {

    switch (type) {
    case GL_FLOAT:
        return "float";
    case GL_FLOAT_VEC2:
        return "vec2";
    case GL_FLOAT_VEC3:
        return "vec3";
    case GL_FLOAT_VEC4:
        return "vec4";
    case GL_DOUBLE:
        return "double";
    case GL_INT:
        return "int";
    case GL_UNSIGNED_INT:
        return "unsigned int";
    case GL_BOOL:
        return "bool";
    case GL_FLOAT_MAT2:
        return "mat2";
    case GL_FLOAT_MAT3:
        return "mat3";
    case GL_FLOAT_MAT4:
        return "mat4";
    default:
        return "?";
    }
}

void GLSLProgram::validate() {
    if (!isLinked()) {
        std::cout << "Program is not linked" << std::endl;
    }

    GLint status;
    glValidateProgram(handle);
    glGetProgramiv(handle, GL_VALIDATE_STATUS, &status);

    if (GL_FALSE == status) {
        // Store log and return false
        int length = 0;
        string logString;

        glGetProgramiv(handle, GL_INFO_LOG_LENGTH, &length);

        if (length > 0) {
            char* c_log = new char[length];
            int written = 0;
            glGetProgramInfoLog(handle, length, &written, c_log);
            logString = c_log;
            delete[] c_log;
        }

        std::cout<<string("Program failed to validate ") + logString<<std::endl;

    }
}

bool GLSLProgram::fileExists(const string& fileName) {
    struct stat info;
    int ret = -1;

    ret = stat(fileName.c_str(), &info);
    return 0 == ret;
}

